import cv2
import numpy as np
from scipy.interpolate import CubicSpline

background = 255-np.zeros((255, 255))
sample_dot = np.array([[[50,50], [100,100], [200,200]]]*7)
img = cv2.imread("lenna.png")
result = img.copy()

def line(x):
    return x

cs_stack = [line]*7

def get_curve(event, x1, y1, flags, param):
    global img1
    if event == cv2.EVENT_LBUTTONDOWN:
        channel = cv2.getTrackbarPos('Channel', 'image')
        control_sample = cv2.getTrackbarPos('Sample Dot', 'image')
        y1 = 255-y1
        sample_dot[channel][1] = np.array([x1,y1])
        sample_dot[channel][0] = np.array([(x1+255)//2 , ((255+y1)//2+(255-x1)//2+y1)//2])
        sample_dot[channel][2] = np.array([(x1+0)//2 , ((0+y1)//2+(0-x1)//2+y1)//2])
        pts = np.vstack((sample_dot[channel],np.array([[0,0],[255,255]])))
        pts[:,0].sort()
        pts[:,1].sort()
        cs_stack[channel] = CubicSpline(pts[:,0], pts[:,1])
        img1 = adjust_h(img, cs_stack[0])
        img1 = adjust_s(img1, cs_stack[1])
        img1 = adjust_v(img1, cs_stack[2])
        img1 = adjust_b(img1, cs_stack[3])
        img1 = adjust_g(img1, cs_stack[4])
        img1 = adjust_r(img1, cs_stack[5])
        img1 = adjust_gamma(img1, cs_stack[6])
        img1 = adjust_dark(img1)
        x = np.arange(255)
        y = 255 - cs_stack[channel](x)
        background = 255-np.zeros((255, 255))
        cv2.polylines(background, np.int32(np.vstack((x,y)).T).reshape((-1,1,2)), True, 0, 2)
        result = img1.copy()
        img1[:,:,2] = img1[:,:,2] * blur /255
        cv2.imshow("image", background)
        cv2.imshow("lomo", img1)

def adjust_h(bgr_img, function):
    hsv_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2HSV)
    hsv_img[:,:,0] = np.array(list(map(lambda x:function(x), hsv_img[:,:,0]/180*255)))/255*180
    result = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)
    return result

def adjust_s(bgr_img, function):
    hsv_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2HSV)
    hsv_img[:,:,1] = np.array(list(map(lambda x:function(x), hsv_img[:,:,1])))
    result = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)
    return result

def adjust_v(bgr_img, function):
    hsv_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2HSV)
    hsv_img[:,:,2] = np.array(list(map(lambda x:function(x), hsv_img[:,:,2])))
    result = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)
    return result

def adjust_b(bgr_img, function):
    result = bgr_img.copy()
    result[:,:,0] = np.array(list(map(lambda x:function(x), bgr_img[:,:,0])))
    return result

def adjust_g(bgr_img, function):
    result = bgr_img.copy()
    result[:,:,1] = np.array(list(map(lambda x:function(x), bgr_img[:,:,1])))
    return result

def adjust_r(bgr_img, function):
    result = bgr_img.copy()
    result[:,:,2] = np.array(list(map(lambda x:function(x), bgr_img[:,:,2])))
    return result

def adjust_gamma(bgr_img, function):
    result = bgr_img.copy()
    result[:,:,0] = np.array(list(map(lambda x:function(x), bgr_img[:,:,0])))
    result[:,:,1] = np.array(list(map(lambda x:function(x), bgr_img[:,:,1])))
    result[:,:,2] = np.array(list(map(lambda x:function(x), bgr_img[:,:,2])))
    return result

def adjust_dark(bgr_img):
    hsv_img = cv2.cvtColor(bgr_img, cv2.COLOR_BGR2HSV)
    hsv_img[:,:,2] = hsv_img[:,:,2] * blur /255
    result = cv2.cvtColor(hsv_img, cv2.COLOR_HSV2BGR)
    return result

def nothing(x):
    pass

mask = np.zeros((img.shape[0], img.shape[1]))
cv2.ellipse(mask, (img.shape[1]//2, img.shape[0]//2), (img.shape[1]//2, img.shape[0]//2), 0, 0, 360, 255, -1)
blur = cv2.GaussianBlur(mask,(301,301),0)

cv2.namedWindow("image")
cv2.resizeWindow("image", 255, 255)
cv2.imshow("image", background)
cv2.setMouseCallback("image", get_curve)
cv2.createTrackbar('Channel', 'image', 0, 6, nothing)
if cv2.waitKey(0) & 0xFF == ord('q'):
    cv2.imwrite('lomo.jpg', img1)
    np.save("lomo", sample_dot)